package com.mz.mattress.client.util;

import java.util.ArrayList;

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

/**
 * https://blog.csdn.net/withme977/article/details/139602904
 * <p>
 * https://juejin.cn/post/7122014462181113887
 */
@Slf4j
public class SseEmitterUtil {

    //userId -> SseEmitter
    private final static Map<Long, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();

    public static SseEmitter connect(Long userId) {

        // 设置超时时间，0表示不过期。默认30S，超时时间未完成会抛出异常：AsyncRequestTimeoutException
        SseEmitter sseEmitter = new SseEmitter(0L);

        // 注册回调
        sseEmitter.onCompletion(completionCallBack(userId));
        sseEmitter.onError(errorCallBack(userId));
        sseEmitter.onTimeout(timeoutCallBack(userId));

        sseEmitterMap.put(userId, sseEmitter);

        log.info("ssw_new_connect,userId={},connect_total={}", userId, sseEmitterMap.size());
        return sseEmitter;
    }

    /**
     * 给制定用户发送消息
     * @param userId     指定用户名
     * @param sseMessage 消息体
     */
    public static void sendMessage(Long userId, String sseMessage) {
        if (sseEmitterMap.containsKey(userId)) {
            try {
                sseEmitterMap.get(userId).send(sseMessage);
                log.info("sse_user={},message={}", userId, sseMessage);
            } catch (IOException e) {
                log.error("sse_error={}", userId, e);
                removeUser(userId);
            }
        } else {
            log.info("sse_user_not_find={},message={}", userId, sseMessage);
        }
    }

    /**
     * 群发消息
     */
    public static void batchSendMessage(String message, List<Long> ids) {
        ids.forEach(userId -> sendMessage(userId, message));
    }

    /**
     * 群发所有人
     */
    public static void batchSendMessage(String message) {
        sseEmitterMap.forEach((k, v) -> {
            try {
                v.send(message, MediaType.APPLICATION_JSON);
            } catch (IOException e) {
                log.error("sse_error={}", k, e);
                removeUser(k);
            }
        });
    }

    /**
     * 移除用户连接
     * @param userId 用户 ID
     */
    public static void removeUser(Long userId) {
        if (sseEmitterMap.containsKey(userId)) {
            sseEmitterMap.get(userId).complete();
            sseEmitterMap.remove(userId);
            log.info("sse_remove={}, remain_connect={}", userId, sseEmitterMap.size());
        }
    }

    /**
     * 获取当前连接信息
     * @return 所有的连接用户
     */
    public static List<Long> getIds() {
        return new ArrayList<>(sseEmitterMap.keySet());
    }

    /**
     * 获取当前的连接数量
     * @return 当前的连接数量
     */
    public static int getUserCount() {
        return sseEmitterMap.size();
    }

    private static Runnable completionCallBack(Long userId) {
        return () -> {
            log.info("sse_end={}", userId);
        };
    }

    private static Runnable timeoutCallBack(Long userId) {
        return () -> {
            log.error("sse_timeout={}", userId);
            removeUser(userId);
        };
    }

    private static Consumer<Throwable> errorCallBack(Long userId) {
        return throwable -> {
            log.error("sse_error={}", userId);
            removeUser(userId);
        };
    }

}
